(cl:in-package #:cleavir-utilities)

(defun depth-first-search-preorder (start-node successor-fun)
  (let ((table (make-hash-table :test #'eq))
        (result '()))
    (labels ((traverse (node)
               (unless (gethash node table)
                 (setf (gethash node table) t)
                 (push node result)
                 (loop for succ in (funcall successor-fun node)
                       do (traverse succ)))))
      (traverse start-node))
    result))

(defun depth-first-search-reverse-postorder (start-node successor-fun)
  (let ((table (make-hash-table :test #'eq))
        (result '()))
    (labels ((traverse (node)
               (unless (gethash node table)
                 (setf (gethash node table) t)
                 (loop for succ in (funcall successor-fun node)
                       do (traverse succ))
                 (push node result))))
      (traverse start-node))
    result))

(defun depth-first-search-postorder (start-node successor-fun)
  (nreverse (depth-first-search-reverse-post-order start-node successor-fun)))

(defun count-nodes (start-node successor-fun)
  (let ((table (make-hash-table :test #'eq)))
    (labels ((traverse (node)
               (unless (gethash node table)
                 (setf (gethash node table) t)
                 (loop for succ in (funcall successor-fun node)
                       do (traverse succ)))))
      (traverse start-node))
    (hash-table-count table)))

;;; Map a function over all nodes.
(defun map-nodes (start-node successor-fun function)
  (let ((table (make-hash-table :test #'eq)))
    (labels ((traverse (node)
               (unless (gethash node table)
                 (setf (gethash node table) t)
                 (funcall function node)
                 (loop for succ in (funcall successor-fun node)
                       do (traverse succ)))))
      (traverse start-node)))
  nil)

(defun predecessor-function (start-node successor-fun)
  (let ((pred-table (make-hash-table :test #'eq)))
    (flet ((successors (node)
             (funcall successor-fun node))
           (predecessors (node)
             (gethash node pred-table))
           ((setf predecessors) (new-predecessors node)
             (setf (gethash node pred-table) new-predecessors)))
      (let ((table (make-hash-table :test #'eq)))
        (labels ((traverse (node)
                   (unless (gethash node table)
                     (setf (gethash node table) t)
                     (loop for succ in (successors node)
                           do (push node (predecessors succ)))
                     (loop for succ in (funcall successor-fun node)
                           do (traverse succ)))))
          (traverse start-node)))
      #'predecessors)))
